import os
import argparse
import csv

import numpy as np


# Gather all numpy array from .npz into one dict
def gather_npz_data(exp_path):
    result_dict = {}
    for file_name in os.listdir(exp_path):
        if ".npz" not in file_name:
            continue
        nd_arrays = np.load(os.path.join(exp_path, file_name))
        for key in nd_arrays.files:
            result_dict[key] = nd_arrays[key]

    return result_dict


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # Granularity : Model-Level / API-Level / Experiment-Level
    parser.add_argument("-s", "--source_path", default="/home/ubuntu/Ascend/results/tf2cann/tensorflow-gpu2",
                        help="The output of source model.")
    parser.add_argument("-f", "--follow_path", default="/home/ubuntu/Ascend/results/tf2cann/om",
                        help="The output of follow-up model.")
    parser.add_argument("-a", "--api_name", default="tf.math.tan", help="The name of API, eg: tf.nn.conv2d")
    parser.add_argument("-o", "--output_path", default="/home/ubuntu/onnx_samples/onnx_transformer/source")
    args = parser.parse_args()
    source_path = os.path.join(args.source_path, args.api_name)
    follow_path = os.path.join(args.follow_path, args.api_name)
    # source_path = args.source_path
    # follow_path = args.follow_path

    result_path_list = []
    # result_path_list = ["20221224"]

    for source_exp_path in os.listdir(source_path):
        if os.path.exists(os.path.join(follow_path, source_exp_path)):
            result_path_list.append(source_exp_path)

    headers = ["name", "result"]
    # input_dict = np.load("/home/ubuntu/onnx_samples/onnx_transformer/source/tf.math.abs_seeds.npz")

    for result_path in result_path_list:
        source_result_dict = gather_npz_data(os.path.join(source_path, result_path))
        follow_result_dict = gather_npz_data(os.path.join(follow_path, result_path))
        # source_result_dict = gather_npz_data(source_path)
        # follow_result_dict = gather_npz_data(follow_path)
        results = []
        source_keys = source_result_dict.keys()
        follow_keys = follow_result_dict.keys()
        all_keys = list(set(source_keys).union(set(follow_keys)))
        tf_error_num = 0
        om_error_num = 0
        passed_num = 0
        failed_num = 0

        # print("[Compare Details]  Violation list:\n")

        for key in all_keys:
            single_result = [key]
            if key not in source_keys:
                single_result.append("TF error")
                tf_error_num += 1
            elif key not in follow_keys:
                single_result.append("OM error")
                om_error_num += 1
            else:
                source_arr = source_result_dict[key]
                follow_arr = follow_result_dict[key]
                if np.allclose(source_arr, follow_arr, atol=0.05, equal_nan=True):
                    single_result.append("Equal")
                    passed_num += 1
                else:
                    single_result.append(source_arr.shape)
                    single_result.append(follow_arr.shape)
                    single_result.append("Not equal")
                    failed_num += 1
            results.append(single_result)

        # print("[Compare] Running %s is over." % args.api_name)
        # print("[Compare] Number of Total Testing Case: %d" % len(all_keys))
        # print("[Compare] Number of Source Error: %d" % tf_error_num)
        # print("[Compare] Number of Follow-up Error: %d" % om_error_num)
        # print("[Compare] Number of Passing: %d, Passing Rate: %.02f" % (passed_num, passed_num / len(all_keys)))
        # print("[Compare] Number of Failed: %d, Failure Rate: %.02f" % (failed_num, failed_num / len(all_keys)))

        with open(args.output_path + "/" + args.api_name + "_" + result_path + ".csv", "w+") as f:
            f_csv = csv.writer(f)
            f_csv.writerow(headers)
            f_csv.writerows(results)
            f_csv.writerow(["tf_error_num", "om_error_num", "passed_num", "failed_num", "total"])
            f_csv.writerow([tf_error_num, om_error_num, passed_num, failed_num, len(all_keys)])


